#include <iostream>
#include <vector>
#include <ctime>
#include <numeric>

#include "tensor.h"
#include "covariance.h"
#include "mtrand.h"
#include "dataset.h"
#include "pdf.h"
#include <vector>
#include <sstream>

#include "TFile.h"
#include "TChain.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TProfile2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TF3.h"
#include "TLorentzVector.h"
#include "TGraphErrors.h"
#include "TGraph2DErrors.h"
#include "make_list.h"

using namespace std;

//initialize random number generator
MTRand_int32 irand(time(0));

const double PI=3.14159265359;
const double min_pt=150.;

bool correction_flag=false;

//counter for number of events
class Counter{
  
public:

  double n;
  double n2;
  double error;
  double tot_n2;
  double tot_n;
  
  
  Counter()
  {
    n=0;
    n2=0;
    error=0;
    tot_n=0;
    tot_n2=0;
  }

  void Add(double weight, double error_=0)
  {
    n+=weight;
    n2+=weight*weight;
    error+=error_;
  }

  
  void Add_lhc(double evt_weight, double frac, double error_=0)
  {
    n+=evt_weight*frac;
    n2+=evt_weight*frac*frac;
    error+=error_;
  }
  
  void Add_tot_n(double tot_err)
  {
    tot_n += tot_err;
    tot_n2 += tot_err*tot_err;
  }

  void Add_tot_n_lhc(double evt_weight, double tot_err)
  {
    tot_n += evt_weight*tot_err;
    tot_n2 += evt_weight*tot_err*tot_err;
  }

  double tot_err() const
  {
    return n*sqrt(tot_n2)/tot_n;
  }

  double nevt() const
  {
    return n;
  }
    
  double err() const
  {
    return sqrt(n2 + error*error);
  }

  double bias() const
  {
    return error;
  }

  double var() const
  {
    return sqrt(n2);
  }

};


class Vec
{
public:
  vector<double> x,y,err;

  Vec(){}

  void push_back(double x_, double y_, double yerr_)
  {
    x.push_back(x_);
    y.push_back(y_);
    err.push_back(yerr_);    
  }

  TGraphErrors get_graph(string name, bool logx=false) const
  {
    vector<double> xprime(x);
    if(logx)
      {
	for(unsigned int i=0; i<xprime.size(); ++i)
	  {
	    xprime[i]=exp10(-2*xprime[i]);
	  }
      }
    TGraphErrors result(xprime.size(), &xprime[0], &y[0], NULL, &err[0]);
    result.SetTitle(name.c_str());
    result.SetName(name.c_str());

    return result;
  }

  //get a 1D graph from pt
  void fill_data(PDF& pdf, double pt, DataSet* data=NULL)
  {
    bool temp= DataSet::b_bias_only;
    DataSet::b_bias_only = true;

    PDF::MC_Generator gen=
      pdf.Generator(make_list<PDF::MC_Flag>
		    (PDF::OUTPUT, PDF::INPUT));

    double ptt=log(pt/min_pt);

    vector<PDF::MC_Event> evt=
      gen.GenerateFull(make_list<double>(ptt));

    for(unsigned int i=0; i<evt.size(); ++i)
      {
	x.push_back(evt[i].value[0]);

	if(correction_flag)
	  y.push_back(evt[i].weight - evt[i].err);

	else
	  y.push_back(evt[i].weight);
      }

    //now compute the bootstrapping error
    err.clear();
    err.resize(y.size(),0);
    
    if(data)
      for(int n=0; n<100; ++n)
      {
	if(n%10==0)
	  cout<<"Bootstrapping: "<<n<<" % done"<<endl;

	//smear the dataset
	DataSet temp_data=data->GenerateDataSet();
	PDF temp_pdf=temp_data.ComputePDF();
	
	PDF::MC_Generator temp_gen=
	  temp_pdf.Generator(make_list<PDF::MC_Flag>
			     (PDF::OUTPUT, PDF::INPUT));
	
	vector<PDF::MC_Event> temp_evt=
	  temp_gen.GenerateFull(make_list<double>(ptt));

	//make sure err is the same size
	assert(err.size()==temp_evt.size());

	for(unsigned int i=0; i<temp_evt.size(); ++i)
	  {
	    double temp_num;

	    if(correction_flag)
	      temp_num=temp_evt[i].weight - temp_evt[i].err - y[i];
	    
	    else
	      temp_num=temp_evt[i].weight - y[i];

	    err[i]+=(temp_num*temp_num);
	  }
      }


    for(unsigned int i=0; i<err.size(); ++i)
      {
	err[i]=sqrt(err[i]/100.);
      }
    
    DataSet::b_bias_only = temp;
  }

};



float DeltaR(float eta, float phi)
{
  if(fabs(phi)> PI)
    phi= PI-fabs(phi);
  return sqrt(eta*eta + phi*phi);
}

int main(int argc, char **argv)
{
  if(argc < 3+1)
    {
      cout<<"Usage: ./madgraph_multijet madgraph_input output.root n_exc [extra smoothing] [correction_flag] [lhc lumi] "
	  <<endl;
      throw;
    }

  if(argc > 3+1)
    {
      cout<<"Extra smoothing factor on "<<atof(argv[4])<<endl;
      DataSet::kernel_smoothing_factor=atof(argv[4]);
    }

  //extra smearing
  bool b_smear_PDF=false;
  bool b_smear_lhc=false;
  double lhc_factor= 1e9;


  if(argc >= 5+1)
    {
      if(atoi(argv[5])==1)
	{
	  cout<<"correction flag enabled (get corrected plot instead)"<<endl;
	  correction_flag=true;
	}
    }

  if(argc >= 6+1)
    {
      cout<<"LHC mode enabled, reading luminosity "
	  <<atof(argv[6])<<"pb^{-1}"<<endl;
      
      b_smear_lhc=true;
      lhc_factor=atof(argv[6])*1e9;
    }


  cout<<"madgraph inputs are read as madgraph_input.?.root ranges from 1 to 12"
      <<endl;
  cout<<"Outputs are automatically stored in PDF_*"<<endl;
  
  //declare variables
  const int maxsize=20;
  Int_t           size;
  Double_t         w;   //[size]
  Float_t         pt[maxsize];   //[size]
  Float_t         eta[maxsize];
  Float_t         phi[maxsize];
  Float_t         m[maxsize];   //[size]
  Float_t         t1[maxsize];
  Float_t         t2[maxsize];
  Float_t         t3[maxsize];

  //exclusive 3jet
  int n_exc=atoi(argv[3]);


  //list of DataSet
  vector<DataSet*> dataset_mpt_njet(4);
  vector<DataSet*> dataset_tau21_njet(4);
  
  DataSet* dataset_mpt_njet_both;
  DataSet* dataset_tau21_njet_both;

  for(int n=0; n<4 && n<n_exc; n++)
    {
      dataset_mpt_njet[n]=new
	DataSet(make_list<double>(0, 0),
		make_list<double>(2.5, 3.0),
		make_list<double>(0.007, 0.007));
  
      dataset_tau21_njet[n]=new
	DataSet(make_list<double>(0, 0),
		make_list<double>(1.0, 3.0),
		make_list<double>(0.006, 0.007));
    }


  dataset_mpt_njet_both=new
    DataSet(make_list<double>(0, 0),
	    make_list<double>(2.5, 3.0),
	    make_list<double>(0.007, 0.007));
  
  dataset_tau21_njet_both=new
    DataSet(make_list<double>(0, 0),
	    make_list<double>(1.0, 3.0),
	    make_list<double>(0.006, 0.007));
      
  //now process madgraph events
  char alist[3]={'A', 'B', 'C'};


  //first get the lhc_prescale
  vector<Counter> lhc_prescale(3, Counter());
  double lhc_prescale_factor[3];
  double lhc_prescale_cut[3]={min_pt, floor(sqrt(280*min_pt)/5)*5, 280};

  //first loop over everything to compute prescale factor
  for(unsigned int k=1; k<=12; k++)
    for(unsigned int ka=0; ka<3; ka++)
      {
	TChain chain_madgraph("Ntuple");
	stringstream sout;
	sout<<argv[1]<<alist[ka]<<"."<<k<<".root";
	chain_madgraph.Add(sout.str().c_str());
	
	//read the input file
	chain_madgraph.SetBranchAddress("size", &size);
	chain_madgraph.SetBranchAddress("w", &w);
	chain_madgraph.SetBranchAddress("pt", pt);
	chain_madgraph.SetBranchAddress("eta", eta);
	chain_madgraph.SetBranchAddress("phi", phi);
	chain_madgraph.SetBranchAddress("m", m);
	chain_madgraph.SetBranchAddress("tau1", t1);
	chain_madgraph.SetBranchAddress("tau2", t2);
	chain_madgraph.SetBranchAddress("tau3", t3);
	
	double rescale=1.0/chain_madgraph.GetEntries();
	
	for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	  {
	    if(i%100000==0)
	      cout<<"\r\033[K"<<floor((i/(chain_madgraph.GetEntries()*12.*3)  +
				       (3*(k-1)+ka)/(12.*3)
				       )*10000)*0.01
		  <<"% pre-processed (for prescale)..."
		  <<flush;
	  

	    chain_madgraph.GetEntry(i);
	    if(size<2) continue;

	    w*=rescale;
	    
	    int real_njet=0;
	    for(int n=0; n<size; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;
		real_njet++;
	      }
	    
	    //if(real_njet!=n_exc) continue;
	    
	    int n_real=0;
	    //limit to two jets
	    for(int n=0; n<size && n_real<n_exc && n_real<2; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;
		
		/*
		double logpt=log(pt[n]/min_pt);
		double mpt=-log10(m[n]/pt[n]);
		double t21=t2[n]/t1[n];
		*/

		//see which prescale bin it is in
		if(n_real==2-1)
		for(unsigned int il=0; il<3; il++)
		  {
		    if(pt[n]> lhc_prescale_cut[il])
		      lhc_prescale[il].Add(w);
		  }

		n_real++;
	      }


	  }
      }
  cout<<"\r\033[K100% pre-processed."<<endl<<flush;

  //now compute all the prescale factors
  const double lhc_Hz=100.;

  for(unsigned int il=0; il<3; il++)
    {
      //the number correspond to cxn in millibarns
      //first compute event rate
      //assume instantenous luminosity = 10^{-34} cm^{-2}
      // = 10 nanobarn^{-1} = 1e7 milibarn^{-1}
      double event_rate= 1e7 * lhc_prescale[il].nevt();

      cout<<"event rate factor for PT bin "<<lhc_prescale_cut[il]
	  <<" is: "<<lhc_prescale[il].nevt()<<endl;

      //make sure event_rate is positive
      assert(event_rate>0);

      //we now want to resctrict it to 200 Hz
      lhc_prescale_factor[il]= lhc_Hz/ event_rate;

      if(lhc_prescale_factor[il]>1.0)
	lhc_prescale_factor[il]=1.0;

      //make sure it is positive
      cout<<"Prescale factor for PT bin "<<lhc_prescale_cut[il]
	  <<" is: "<<lhc_prescale_factor[il]<<endl;
      assert(lhc_prescale_factor[il]>0);      
    }

  


  for(unsigned int k=1; k<=12; k++)
    for(unsigned int ka=0; ka<3; ka++)
    {
      TChain chain_madgraph("Ntuple");
      stringstream sout;
      sout<<argv[1]<<alist[ka]<<"."<<k<<".root";
      chain_madgraph.Add(sout.str().c_str());

      //read the input file
      chain_madgraph.SetBranchAddress("size", &size);
      chain_madgraph.SetBranchAddress("w", &w);
      chain_madgraph.SetBranchAddress("pt", pt);
      chain_madgraph.SetBranchAddress("eta", eta);
      chain_madgraph.SetBranchAddress("phi", phi);
      chain_madgraph.SetBranchAddress("m", m);
      chain_madgraph.SetBranchAddress("tau1", t1);
      chain_madgraph.SetBranchAddress("tau2", t2);
      chain_madgraph.SetBranchAddress("tau3", t3);

      double rescale=1.0/chain_madgraph.GetEntries();

      for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	{
	  if(i%1000==0)
	    cout<<"\r\033[K"<<(i*100000/(chain_madgraph.GetEntries()*12*3)  +
			       (3*(k-1)+ka)*100000/(12*3)
			       )*0.001
		<<"% trained..."
		<<flush;

	  
	  
	  chain_madgraph.GetEntry(i);
	  if(size<2) continue;
	 
	  w*=rescale;

	  int real_njet=0;
	  for(int n=0; n<size; ++n)
	    {
	      if(t1[n] < 1e-5) continue;
	      if(t2[n] < 1e-5) continue;
	      if(pt[n] < min_pt) continue;
	      if(m[n]<20) continue;
	      real_njet++;
	    }

	  if(real_njet!=n_exc) continue;

	     /*
	  if(size>n_exc && pt[n_exc] > 50.0)
	    continue;
	     */

	  int n_real=0;

	  //get the prescale factor;
	  for(int n=0; n<size && n_real<n_exc; ++n)
	    {
	      if(t1[n] < 1e-5) continue;
	      if(t2[n] < 1e-5)  continue;
	      if(pt[n] < min_pt) continue;
	      if(m[n]<20) continue;

	      n_real++;
	      //if it is the second jet
	      //get the pre-scale
	      if(n_real==2)
		for(unsigned int il=0; il<3; il++)
		  {
		    //find prescale bin
		    if(pt[n]>= lhc_prescale_cut[il] && 
		       (il>=2 || pt[n]< lhc_prescale_cut[il+1]))
		      {
			w*=lhc_prescale_factor[il];
			break;
		      }
		  }
	    }

	  //reset n_real counter
	  n_real=0;

	  for(int n=0; n<size && n_real<n_exc; ++n)
	    {
	      if(t1[n] < 1e-5) continue;
	      if(t2[n] < 1e-5)  continue;
	      if(pt[n] < min_pt) continue;
	      if(m[n]<20) continue;
	      
	      double logpt=log(pt[n]/min_pt);
	      double mpt=-log10(m[n]/pt[n]);
	      double t21=t2[n]/t1[n];
	      
	      dataset_mpt_njet[n_real]->Fill(make_list<double>(mpt,logpt),
					     w);
	      dataset_tau21_njet[n_real]->Fill(make_list<double>(t21,logpt),
					       w);

	      if(n_real<2)
		{
		  dataset_mpt_njet_both->Fill(make_list<double>(mpt,logpt),
					      w);
		  dataset_tau21_njet_both->Fill(make_list<double>(t21,logpt),
					       w);
		}

	      ++n_real;


	    }
	}
      
    }
  cout<<"\r\033[K100% trained."<<endl<<flush;


  //get PDF list
  PDF PDF_mpt_njet[4];
  PDF PDF_tau21_njet[4]; 

  PDF PDF_mpt_njet_both;
  PDF PDF_tau21_njet_both; 

  DataSet* mpt_dataset[4];
  DataSet* tau21_dataset[4];

  DataSet* mpt_dataset_both;
  DataSet* tau21_dataset_both;

  //for the njet cases
  for(int n=0; n<dataset_mpt_njet.size() && n<n_exc; ++n)
    {
      //get the smeared pdf
      if(b_smear_lhc)
	{
	  DataSet tt= dataset_mpt_njet[n]->GenerateDataSet(lhc_factor);
	  PDF_mpt_njet[n]=tt.ComputePDF();
	  mpt_dataset[n]=new DataSet(tt);

	  DataSet ttt=dataset_tau21_njet[n]->GenerateDataSet(lhc_factor);
	  PDF_tau21_njet[n]=ttt.ComputePDF();      	    
	  tau21_dataset[n]=new DataSet(ttt);

	  delete dataset_mpt_njet[n];
	  delete dataset_tau21_njet[n];
	}  

      else
	{
	  PDF_mpt_njet[n]=dataset_mpt_njet[n]->ComputePDF();
	  PDF_tau21_njet[n]=dataset_tau21_njet[n]->ComputePDF();      

	  mpt_dataset[n]=dataset_mpt_njet[n];
	  tau21_dataset[n]=dataset_tau21_njet[n];
	}

      //now delete the dataset
      /*
      delete dataset_mpt_njet[n];
      delete dataset_tau21_njet[n];
      */
    }

  //templates for both jets
  if(b_smear_lhc)
    {
      DataSet tt= dataset_mpt_njet_both->GenerateDataSet(lhc_factor);
      PDF_mpt_njet_both=tt.ComputePDF();
      mpt_dataset_both=new DataSet(tt);      
      
      DataSet ttt=dataset_tau21_njet_both->GenerateDataSet(lhc_factor);
      PDF_tau21_njet_both=ttt.ComputePDF();      	    
      tau21_dataset_both=new DataSet(ttt);      

      //now delete the dataset
      delete dataset_mpt_njet_both;
      delete dataset_tau21_njet_both;
    }

  else
    {
      PDF_mpt_njet_both=dataset_mpt_njet_both->ComputePDF();
      PDF_tau21_njet_both=dataset_tau21_njet_both->ComputePDF();      

      mpt_dataset_both=dataset_mpt_njet_both;
      tau21_dataset_both=dataset_tau21_njet_both;
    }
  
    
  
  //get output
  TFile out(argv[2], "RECREATE");
  map<string, Vec> graph_list;

  for(int i=400; i<=600; i+=100)
    {
      stringstream sout;
      
      //do the njet case
      for(int n=0; n<4 && n<n_exc; ++n)
	{
	  sout.str("");
	  sout<<"t21_"<<n<<"jet_"<<i;
	  graph_list[sout.str()].fill_data(PDF_tau21_njet[n], i, tau21_dataset[n]);
	  graph_list[sout.str()].get_graph(sout.str()).Write();

	  sout.str("");
	  sout<<"mpt_"<<n<<"jet_"<<i;
	  graph_list[sout.str()].fill_data(PDF_mpt_njet[n], i, mpt_dataset[n]);
	  graph_list[sout.str()].get_graph(sout.str(), true).Write();	  
	}

      //do both jets templates
      sout.str("");
      sout<<"mpt_"<<"both"<<"jet_"<<i;
      graph_list[sout.str()].fill_data(PDF_mpt_njet_both, i, mpt_dataset_both);
      graph_list[sout.str()].get_graph(sout.str(), true).Write();     


      sout.str("");
      sout<<"t21_"<<"both"<<"jet_"<<i;
      graph_list[sout.str()].fill_data(PDF_tau21_njet_both, i, tau21_dataset_both);
      graph_list[sout.str()].get_graph(sout.str()).Write();     

    }


  //write histograms
  /*
  for(map<string, TH1D*>::iterator it=hist.begin();
      it!=hist.end(); ++it)
    {
      it->second->Write();
    }
  */

  cout<<"Program ends."<<endl;
  return 1;
}
  
